#---------------------# 
# Read stratification #
#---------------------# 

#--------------------
# Import Modules
#--------------------

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.axes import Axes
from functions import PT  
from mpl_toolkits.axes_grid1 import ImageGrid
from scipy.integrate import cumtrapz 

#--------------------
# Constants 
#--------------------

G = 6.674e-11    # [G] = [N(m/kg)^2]
R = 13732.   
rs = 6.96e8      # 4.55Gyr

#-------------------------------------------------------
# Read Environmental Data with Polytropic aproximation  
#-------------------------------------------------------

#inputfile1 ='/home/zaire/mfdyn/trunk/sms/data.txt'
inputfile1 ='/home/guerrero/mfdyn/trunk/sms/data.sun'

f_env = open(inputfile1, 'r')
files_env = np.genfromtxt(f_env, usecols = (0, 1, 2, 3))
r_env = files_env[:,0]     # [r] = [m]
T_env = files_env[:,1]     # [T_e] = [K]                  
rho_env = files_env[:,2]   # [rho_e] = [Kg/m^3]  
pt_env = files_env[:,3]
f_env.close()

r = r_env
 
#r = np.linspace(0.6*rs, 0.95*rs, 64) 
rr = r/rs

#-------------------------------------------------------
# Read Environmental Data with Polytropic aproximation
#-------------------------------------------------------

inputfile1 ='/home/guerrero/mfdyn/trunk/sms/data.sun'

f_isen = open(inputfile1, 'r')
files_isen = np.genfromtxt(f_isen, usecols = (0, 1, 2, 3))
r_isen = files_isen[:,0]     # [r] = [m]
T_isen = files_isen[:,1]     # [T] = [K]                  
rho_isen = files_isen[:,2]   # [rho] = [Kg/m^3]  
pt_isen = files_isen[:,3]
f_isen.close()

#------------------------------------------
# Read Data Stelar Model
#------------------------------------------

#inputfile ='/home/zaire/eumag/stratification_0.61Msun_6.66Myr'
inputfile ='/home/guerrero/Dropbox/bonnie/stratification/estratification_1Msun_4.55Gyr.out'

f = open(inputfile, 'r')
files = np.genfromtxt( f, usecols = (0,1,2,3,6))
r_sm = 1e-2*files[:,0]  # [r] = [m]
P = files[:,1]  
T = files[:,2]         # [T] = [K]
rho = 1e3*files[:,3]   # [rho] = [Kg/m^3]
dlntdlnp = files[:,4]
f.close()

M = 4 * np.pi * cumtrapz(rho*r_sm**2, r_sm, initial = 0)

## These arrays have points for the same radius used in polytropic model:

P = np.interp(r, r_sm, P)
T = np.interp(r, r_sm, T)
rho = np.interp(r, r_sm, rho)
M = np.interp(r, r_sm, M)
dlntdlnp = np.interp(r, r_sm, dlntdlnp)

g = G * M / r**2
Tb = T[0]
rhob = rho[0]
pt = T*(Tb*rhob/(T*rho))**(2./5.)
pt1 = T*(P[0]/P)**(2./5.)

#--------------------------------------------
# Polynomial approximation for Gravity 
#--------------------------------------------

## FOR STELLAR MODEL:

coefficients = np.polyfit(r, g, 3)
polynomial = np.poly1d(coefficients)
g_p = polynomial(r)

#print polynomial
print 'g = gb*(', polynomial/g_p[0], ')'
#print 'rho =', rhob
#print 'T =', Tb
#print 'gb =', g_p[0]

#-----------------------------------------
# Plot
#-----------------------------------------

font = {'family' : 'serif',
        'color'  : 'black',
        'fontweight' : 'normal',
        'fontsize'   : 15,
        'horizontalalignment' : 'center',
        'verticalalignment': 'center',
       }

font1 = {'family' : 'serif'}

plt.rc('text', usetex=True)

plt.plot(rr, g, 'm', rr, g_p, 'b--', linewidth = 1.5)
plt.legend(('Stellar Model', 'Polynomial Fit'), loc = 9, prop = font1, fontsize = 15)
plt.xscale('linear')
plt.yscale('linear')
plt.xlim(rr.min(),rr.max())
plt.ylim(g.min(),1.05*g.max())
#plt.xticks([0.1, 0.2, 0.3, 0.4,0.5,0.6,0.7,0.8,0.9])
plt.ticklabel_format(style = 'scientific', scilimits = (0,3), axis = 'both', pad = 10)
plt.tick_params(axis = 'both', direction = 'inout', labelsize = 15, colors = 'k')
plt.xlabel(r'$r [R_{\rm s}]$', fontdict = font, labelpad=10)
plt.ylabel(r'$g \left [m/s^{2} \right ]$', fontdict = font, labelpad=10)
plt.margins(0.1,0.1)       
plt.show()


plt.plot(rr, T, 'm', rr, T_env, 'b--', rr, T_isen, 'r-.', linewidth = 1.5)
plt.legend((r'$T_{\rm sm}$', r'$T_{\rm e}$', r'$T_{\rm s}$'), loc = 3, prop = font1, fontsize = 15)
plt.xscale('linear')
plt.yscale('linear')
plt.xlim(rr.min(),rr.max())
#plt.ylim(T.min(), 1.01*T_env.max())
#plt.xticks([0.1, 0.2, 0.3, 0.4,0.5,0.6,0.7,0.8,0.9])
plt.ticklabel_format(style = 'scientific', scilimits = (0,3), axis = 'both', pad = 10)
plt.tick_params(axis = 'both', direction = 'inout', labelsize = 15, colors = 'k')
plt.xlabel(r'$r [R_{\rm s}]$', fontdict = font, labelpad=10)
plt.ylabel(r'$T[K]$', fontdict = font, labelpad=10)
plt.margins(0.8,0.8)       
plt.show()


plt.plot(rr, rho, 'm', rr, rho_env, 'b--', rr, rho_isen, 'r-.', linewidth = 1.5)
plt.legend((r'$\rho_{\rm sm}$', r'$\rho_{\rm e}$', r'$\rho_{\rm s}$'), loc = 3, prop = font1, fontsize = 15)
plt.xscale('linear')
plt.yscale('linear')
plt.xlim(rr.min(),rr.max())
#plt.ylim(rho.min(), 1.01*rho_env.max())
#plt.xticks([0.1, 0.2, 0.3, 0.4,0.5,0.6,0.7,0.8,0.9])
plt.ticklabel_format(style = 'scientific', scilimits = (0,3), axis = 'both', pad = 10)
plt.tick_params(axis = 'both', direction = 'inout', labelsize = 15, colors = 'k')
plt.xlabel(r'$r [R_{\rm s}]$', fontdict = font, labelpad=10)
plt.ylabel(r'$\rho \left [Kg/m^{3} \right ]$', fontdict = font, labelpad=10)
plt.margins(0.3,0.3)       
plt.show()
